from mpi4py import MPI
import numpy as np

import ufl
from basix.ufl import element, mixed_element
from dolfinx import default_real_type, fem, la
from dolfinx.fem import (
    Constant,
    Function,
    dirichletbc,
    extract_function_spaces,
    form,
    functionspace,
    locate_dofs_topological,
)
from dolfinx.fem.petsc import assemble_matrix_block, assemble_vector_block
from dolfinx.io import XDMFFile
from dolfinx.mesh import CellType, create_rectangle, locate_entities_boundary
from ufl import div, dx, grad, inner

# PETSc
# Demo文件给出了很多种情况的求解方案: 迭代求解 直接求解...
try:
    from petsc4py import PETSc

    import dolfinx

    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
    from petsc4py.PETSc import ScalarType
except ModuleNotFoundError:
    print("This demo requires petsc4py.")
    exit(0)


# 定义
def true_velocity(x):
    u1 = 20 * x[0] ** 2 * (x[0] - 1) ** 2 * x[1] * (x[1] - 1) * (2 * x[1] - 1)
    u2 = -20 * x[0] * (x[0] - 1) * (2 * x[0] - 1) * x[1] ** 2 * (x[1] - 1) ** 2
    return [u1, u2]


def true_pressure(x):
    return 10 * (2 * x[0] - 1) * (2 * x[1] - 1)


def D(u):
    return 0.5 * (ufl.nabla_grad(u) + ufl.nabla_grad(u).T)


def getStiffnessMatrix(A: PETSc.Mat):
    """
    获取刚度矩阵的值(np.array)
    """
    m, n = A.getSize()
    matrix_data = np.zeros((m, n))

    # 获取矩阵的所有非零元素
    row, col, val = A.getValuesCSR()

    # 将CSR格式的数据转换为密集矩阵
    for i in range(m):
        for j in range(row[i], row[i + 1]):
            matrix_data[i, col[j]] = val[j]

    return matrix_data
    

nu = 1  # 粘度系数
idx = int(input("设置剖分(可选0, 1, 2, 3) = "))  # 剖分索引
n_list = [[16, 16],
          [32, 32],
          [64, 64],
          [128, 128],
          [4, 4],
          ]  # 剖分列表
print(f"{n_list[idx]}剖分...")
msh = create_rectangle(
    MPI.COMM_WORLD, [np.array([0, 0]), np.array(
        [1, 1])], n_list[idx], CellType.triangle
)

# Define Function Space: u1, u2 in V; p in Q
P2 = element(
    "Lagrange", msh.basix_cell(), degree=2, shape=(msh.geometry.dim,), dtype=default_real_type
)
P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
V, Q = functionspace(msh, P2), functionspace(msh, P1)

# 标记边界 单位正方形
facets = locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(
        x[0], 1.0) | np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0),
)

dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)


# 定义边界条件 g
def g(x):
    values = np.zeros((msh.geometry.dim, x.shape[1]))
    u1_expr = 20 * x[0] ** 2 * (x[0] - 1) ** 2 * \
        x[1] * (x[1] - 1) * (2 * x[1] - 1)
    u2_expr = -20 * x[0] * (x[0] - 1) * (2 * x[0] - 1) * \
        x[1] ** 2 * (x[1] - 1) ** 2
    values[0] = u1_expr
    values[1] = u2_expr
    return values


bc_value = Function(V)
bc_value.interpolate(g)
bc = dirichletbc(bc_value, dofs)

# 定义体积力 f
x = ufl.SpatialCoordinate(msh)
f1 = - 240.0 * nu * x[0] ** 4 * x[1] + 120.0 * nu * x[0] ** 4 + 480.0 * nu * x[0] ** 3 * x[1] - 240.0 * nu * x[0] ** 3 - 480.0 * nu * x[0] ** 2 * x[1] ** 3 + 720.0 * nu * x[0] ** 2 * x[1] ** 2 - 480 * nu * x[0] ** 2 * \
    x[1] + 120.0 * nu * x[0] ** 2 + 480.0 * nu * x[0] * x[1] ** 3 - 720.0 * nu * x[0] * x[1] ** 2 + 240.0 * \
    nu * x[0] * x[1] - 80.0 * nu * x[1] ** 3 + 120.0 * \
    nu * x[1] ** 2 - 40.0 * nu * x[1] + 40 * x[1] - 20
f2 = 480.0 * nu * x[0] ** 3 * x[1] ** 2 - 480.0 * nu * x[0] ** 3 * x[1] + 80.0 * nu * x[0] ** 3 - 720.0 * nu * x[0] ** 2 * x[1] ** 2 + 720.0 * nu * x[0] ** 2 * x[1] - 120.0 * nu * x[0] ** 2 + 240.0 * nu * x[0] * \
    x[1] ** 4 - 480.0 * nu * x[0] * x[1] ** 3 + 480 * nu * x[0] * x[1] ** 2 - 240.0 * nu * x[0] * x[1] + 40.0 * \
    nu * x[0] - 120.0 * nu * x[1] ** 4 + 240.0 * nu * \
    x[1] ** 3 - 120.0 * nu * x[1] ** 2 + 40 * x[0] - 20
f = ufl.as_vector([f1, f2])

# 定义变分问题 (关键修正：修正变分形式中的符号错误) 
(u, p) = ufl.TrialFunction(V), ufl.TrialFunction(Q)
(v, q) = ufl.TestFunction(V), ufl.TestFunction(Q)

# 注意这里不能用LinearSolve 只能解决a(u,v) = L(v)的问题 本问题明显更复杂
a = form([[2 * nu * inner(D(u), D(v)) * dx, -inner(p, div(v)) * dx],
          [inner(div(u), q) * dx, None]])
L = form([inner(f, v) * dx, inner(Constant(msh, PETSc.ScalarType(0)), q) * dx])

a_p11 = form(inner(p, q) * dx)
a_p = [[a[0][0], None], [None, a_p11]]


def block_operators():
    """Return block operators and block RHS vector for the Stokes
    problem"""

    # Assembler matrix operator, preconditioner and RHS vector into
    # single objects but preserving block structure
    A = assemble_matrix_block(a, bcs=[bc])
    A.assemble()
    P = assemble_matrix_block(a_p, bcs=[bc])
    P.assemble()
    # 对右端项进行组装
    b = assemble_vector_block(L, a, bcs=[bc])

    # Set the nullspace for pressure (since pressure is determined only
    # up to a constant)
    null_vec = A.createVecLeft()
    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
    # print(f"{V.dofmap.index_map.size_local}, {V.dofmap.index_map_bs}")
    null_vec.array[offset:] = 1.0
    null_vec.normalize()
    nsp = PETSc.NullSpace().create(vectors=[null_vec])
    assert nsp.test(A)
    A.setNullSpace(nsp)

    return A, P, b


def block_direct_solver():
    """Solve the Stokes problem using blocked matrices and a direct
    solver.

    """

    # Assembler the block operator and RHS vector
    A, _, b = block_operators()  # 这里没有使用预条件子矩阵

    # Create a solver
    ksp = PETSc.KSP().create(msh.comm)
    ksp.setOperators(A)
    ksp.setType("preonly")  # 直接求解

    # MUMPS是并行直接求解器
    # Set the solver type to MUMPS (LU solver) and configure MUMPS to
    # handle pressure nullspace
    pc = ksp.getPC()
    pc.setType("lu")
    sys = PETSc.Sys()  # type: ignore
    use_superlu = PETSc.IntType == np.int64
    if sys.hasExternalPackage("mumps") and not use_superlu:
        pc.setFactorSolverType("mumps")
        pc.setFactorSetUpSolverType()
        # icntl=24 允许处理奇异矩阵 =25禁用错误检测
        pc.getFactorMatrix().setMumpsIcntl(icntl=24, ival=1)
        pc.getFactorMatrix().setMumpsIcntl(icntl=25, ival=0)
    else:
        pc.setFactorSolverType("superlu_dist")

    # Create a block vector (x) to store the full solution, and solve
    x = A.createVecLeft()  # 存储解的块向量（速度和压力）
    stiff_mat = getStiffnessMatrix(A=A)  # 刚度矩阵太大会oom(n=128)
    # print(np.linalg.cond(stiff_mat))
    load_vec = b.getArray()  # 载荷向量可以直接获取
    np.save(f"./data/stiff_mat{n_list[idx][0]}", stiff_mat)
    np.save(f"./data/load_vec{n_list[idx][0]}", load_vec)
    print(f"Save StiffMat and LoadVec(n={n_list[idx][0]}) to ./data")
    ksp.solve(b, x)
    # xx = np.linalg.solve(stiff_mat, load_vec)  # 验证求解器的正确性

    # Create Functions and scatter x solution
    u, p = Function(V), Function(Q)
    offset = V.dofmap.index_map.size_local * V.dofmap.index_map_bs
    u.x.array[:offset] = x.array_r[:offset]  # 取出速度和压力的解
    p.x.array[: (len(x.array_r) - offset)] = x.array_r[offset:]
 
    # 计算误差
    u_true = Function(V)  # 实例化空间
    u_true.interpolate(true_velocity)
    p_true = Function(Q)
    p_true.interpolate(true_pressure)

    u_error = Function(V)
    u_error.x.array[:] = u.x.array - u_true.x.array
    p_error = Function(Q)
    p_error.x.array[:] = p.x.array - p_true.x.array

    # Compute norms of the u and p vectors
    u_comm = u.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    u_h1_error = form((grad(u_error)**2 + u_error**2) * ufl.dx)
    u_l2_error = form(u_error**2 * ufl.dx)
    u_h1_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_h1_error), MPI.SUM))
    u_l2_E = np.sqrt(u_comm.allreduce(
        fem.assemble_scalar(u_l2_error), MPI.SUM))  # 并行求和 返回float

    p_comm = p.function_space.mesh.comm
    # type = dolfinx.fem.forms.Form
    p_h1_error = form((grad(p_error)**2 + p_error**2) * ufl.dx)
    p_l2_error = form(p_error**2 * ufl.dx)
    p_h1_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_h1_error), MPI.SUM))
    p_l2_E = np.sqrt(p_comm.allreduce(
        fem.assemble_scalar(p_l2_error), MPI.SUM))  # 并行求和 返回float

    # Linf 比较节点值 离散后区域无法计算理论最大值
    u_linf_E, p_linf_E = la.norm(
        u_error.x, la.Norm.linf), la.norm(p_error.x, la.Norm.linf)
    # 超级重要 是我心目的官方文档
    # https://jsdokken.com/dolfinx-tutorial/chapter4/convergence.html#computing-convergence-rates

    if MPI.COMM_WORLD.rank == 0:
        print(
            f"H1 norm for velocity: {u_h1_E}, pressure: {p_h1_E}.")
        print(
            f"L2 norm for velocity: {u_l2_E}, pressure: {p_l2_E}.")
        print(
            f"Linf norm for velocity: {u_linf_E}, pressure: {p_linf_E}.")
    
    # 绘图
    
    with XDMFFile(MPI.COMM_WORLD, f"figures/velocity-N={n_list[idx][0]}.xdmf", "w") as ufile_xdmf:
        u.x.scatter_forward()  # u是数值结果 type= dolfinx.fem.Function
        P1 = element(
            "Lagrange", msh.basix_cell(), 1, shape=(msh.geometry.dim,), dtype=default_real_type
        )  # 画图只需要P1元即可 效率更高
        u1 = Function(functionspace(msh, P1))
        u1.interpolate(u)  # 根据计算结果进行插值
        ufile_xdmf.write_mesh(msh)
        ufile_xdmf.write_function(u1)

    with XDMFFile(MPI.COMM_WORLD, f"figures/pressure-N={n_list[idx][0]}.xdmf", "w") as pfile_xdmf:
        p.x.scatter_forward()
        pfile_xdmf.write_mesh(msh)
        pfile_xdmf.write_function(p)

    print(f"N={n_list[idx][0]}, 已保存绘图结果至figures文件夹...\n")

block_direct_solver()
